import torch.optim as optim
import os.path as osp

from torch_geometric.data import DataLoader
from lib.models.gnn import GNN
from lib.utils.get_config import get_cfg
from lib.evaluation_metrics.metrics import get_acc
from lib.loss.loss import get_loss
# read config file
cfg = get_cfg(osp.join(osp.abspath("./"), 'configs/config.yaml'))
DEVICE = cfg['DEVICE']


def train(dataset, task, hidden_num, hidden_dim, lr=0.01):
    trigger = 0
    prev_acc = 0.0
    validation_loss, validation_acc, test_acc, test_loss = 0, 0, 0, 0
    if task == 'graph':
        data_size = len(dataset)
        # TODO:cross validation
        # split data  into traning, validation and testing set according to 3:1:1 ratio
        loader = DataLoader(
            dataset[:int(data_size * (1-cfg['PERCENT']*2))], batch_size=cfg['BATCH_SIZE'])
        validation_loader = DataLoader(
            dataset[int(data_size * (1-cfg['PERCENT']*2)):int(data_size * (1-cfg['PERCENT']))], batch_size=cfg['BATCH_SIZE'], shuffle=False)
        test_loader = DataLoader(
            dataset[int(data_size * (1-cfg['PERCENT'])):], batch_size=cfg['BATCH_SIZE'], shuffle=False)
    else:
        test_loader = DataLoader(
            dataset, batch_size=cfg['BATCH_SIZE'], shuffle=cfg['SHUFFLE'])
        loader = test_loader
        validation_loader = test_loader

    # build model
    model = GNN(max(dataset.num_node_features, 1), dataset.num_classes,
                task=task, hidden_num=hidden_num, hidden_dim=hidden_dim)
    model = model.to(DEVICE)
    opt = optim.Adam(model.parameters(), lr)

    # train for EPOCHS epochs
    for _ in range(cfg['EPOCHS']):
        training_loss = 0
        validation_loss = 0

        model.train()
        for batch in loader:
            opt.zero_grad()
            # get the prediction and the excpeted label
            batch = batch.to(DEVICE)
            embedding, pred = model(batch)
            label = batch.y
            if task == 'node':
                pred = pred[batch.train_mask]
                label = label[batch.train_mask]
            loss = model.loss(pred, label)
            loss.backward()
            opt.step()
            training_loss += loss.item() * batch.num_graphs

        training_loss /= len(loader.dataset)

        validation_loss = get_loss(validation_loader, opt, model)
        validation_acc = get_acc(validation_loader, model, True)

        if validation_acc <= prev_acc:
            trigger += 1

        if trigger >= cfg['PATIENCE_TRESHOLD']:
            test_acc = get_acc(test_loader, model, False)
            test_loss = get_loss(test_loader, opt, model)
            break
        prev_acc = validation_acc

    #  TODO:把标红的改为列表.append()形式，最后sum
    return model, validation_loss, validation_acc, test_acc, test_loss
